/*
 * JGO2RDF creates a RDF representation of GAF2 files, which are automatically sourced from GO ftp site.
 * Copyright (C) 2012 Matthew M. Hindle & Artem Lysenko
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package uk.co.jgo2rdf.goa.gp2protein;

import com.csvreader.CsvReader;
import org.apache.commons.net.ftp.FTP;
import org.apache.commons.net.ftp.FTPClient;
import org.apache.log4j.Logger;
import org.apache.lucene.analysis.KeywordAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.*;
import org.apache.lucene.store.NIOFSDirectory;
import org.apache.lucene.util.Version;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.nio.charset.Charset;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.zip.GZIPInputStream;

/**
 * @author mhindle
 *         Date: 06/02/12
 *         Time: 17:39
 */
public class ParseGP2XRefFile {

    private static final Logger log = Logger.getLogger(ParseGP2XRefFile.class);

    private static final String GP = "gp";
    private static final String UNIPROT = "uniprot";

    private Set<String> databaseFrom = new HashSet<String>(); //stores all database ids
    private NIOFSDirectory directory;
    private String indexDir;

    private IndexSearcher indexSearcher; //is automatically init after parsing


    /**
     * @param indexDir  directory of lucene index
     * @param loadIndex if true it will load the existing index at the dir else it is deleted
     */
    public ParseGP2XRefFile(String indexDir, boolean loadIndex) throws IOException {
        this.indexDir = indexDir;

        File f = new File(indexDir);
        if (!loadIndex) { //this is the fist run
            f.delete();
            f.mkdirs();
            f.mkdir();
        }

        directory = new NIOFSDirectory(f);

        if (loadIndex) {
            indexSearcher = new IndexSearcher(IndexReader.open(directory, true), null);
        }
    }

    public Set<String> getDatabasesFrom() {
        return databaseFrom;
    }

    /**
     * Parses a pg2protein file. You may call this successive times to create a single index
     *
     * @param url            the location of the pg2protein file (can be gziped and from ftp)
     * @param fromAccessions the translation from accessions to exclusively include
     * @throws IOException
     */
    public void parseFromURL(URL url, Collection<String> fromAccessions) throws IOException {
        if (url.getProtocol().equalsIgnoreCase("ftp")) {
            log.info("Fetch " + url.getFile() + " from " + url.getHost());

            FTPClient ftp = new FTPClient();
            try {
                ftp.connect(url.getHost());
                ftp.login("anonymous", "anonymous");
                ftp.setFileTransferMode(FTP.STREAM_TRANSFER_MODE);
                ftp.setFileType(FTP.BINARY_FILE_TYPE);
                if (url.getFile().endsWith(".gz")) {
                    parseFromStream(new GZIPInputStream(
                            ftp.retrieveFileStream(url.getFile())), fromAccessions);
                } else {
                    parseFromStream(ftp.retrieveFileStream(url.getFile()), fromAccessions);
                }
            } finally {
                ftp.disconnect();
            }
        } else {
            parseFromStream(url.openStream(), fromAccessions);
        }
    }

    /**
     * Parses a pg2protein file. You may call this successive times to create a single index
     *
     * @param is             input stream to a gp2protein file
     * @param fromAccessions the translation from accessions to exclusively include
     * @throws IOException
     */
    public void parseFromStream(InputStream is, Collection<String> fromAccessions) throws IOException {

        if (indexSearcher != null)
            indexSearcher.close();

        IndexWriterConfig config = new IndexWriterConfig(Version.LUCENE_30,
                new KeywordAnalyzer());
        IndexWriter writer = new IndexWriter(directory, config);

        CsvReader csvReader = new CsvReader(is,
                '\t', Charset.defaultCharset());
        try {
            csvReader.setComment('!');
            csvReader.setUseComments(true);
            csvReader.setTrimWhitespace(false);
            csvReader.setHeaders(new String[]{GP, UNIPROT});

            int i = 0;
            while (csvReader.readRecord()) {
                if (i % 10000 == 0)
                    System.out.println("Processed " + i + " records");
                i++;

                String gp_raw = csvReader.get(GP);
                if (gp_raw == null || gp_raw.length() == 0) continue;

                String[] values = gp_raw.split(":");
                String gp_database = values[0];
                String gp_value = values[1];

                String target_raw = csvReader.get(UNIPROT);
                if (target_raw == null || target_raw.length() == 0) continue;

                String[] xrefs = target_raw.split("|");

                for (String xref : xrefs) {
                    String[] pair = xref.split(":");
                    String t_database = values[0];
                    String t_value = values[1];

                    if (fromAccessions == null || fromAccessions.contains(gp_value) || fromAccessions.contains(gp_value.toUpperCase())) {
                        databaseFrom.add(gp_database.toUpperCase());
                        Document doc = new Document();
                        doc.add(new Field("gp_database",
                                gp_database,
                                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                        doc.add(new Field("gp_value",
                                gp_value,
                                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                        doc.add(new Field("t_database",
                                gp_database,
                                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                        doc.add(new Field("t_value",
                                t_value,
                                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                        writer.addDocument(doc);

                        int version = gp_value.lastIndexOf('.');
                        if (version > 0) {
                            Document docuv = new Document();
                            docuv.add(new Field("gp_database",
                                    gp_database,
                                    Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                            docuv.add(new Field("gp_value",
                                    gp_value.substring(0, version),
                                    Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                            doc.add(new Field("t_database",
                                    gp_database,
                                    Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                            docuv.add(new Field("t_value",
                                    t_value,
                                    Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                            writer.addDocument(docuv);
                        }
                    }
                }
            }
        } finally {
            csvReader.close();
            is.close();
            writer.prepareCommit();
            writer.commit();
            writer.close(true);
        }
        indexSearcher = new IndexSearcher(IndexReader.open(directory, true), null);
    }

    /**
     * Searches for Uniprot translations of an accessions
     *
     * @param accession query accession of unknown database
     * @return uniprot accessions
     * @throws IOException
     */
    public Set<String> getUniprotAccession(String accession) throws IOException {
        return getAccession(null, accession, "UniProtKB");
    }

    /**
     * Searches for Uniprot translations of an accessions
     *
     * @param database  the database the accession query originates from
     * @param accession the query accession
     * @return uniprot accessions
     * @throws IOException
     */
    public Set<String> getUniprotAccession(String database, String accession) throws IOException {
        return getAccession(database, accession, "UniProtKB");
    }

    /**
     * @param database       the database the accession query originates from
     * @param accession      the query accession
     * @param targetdatabase the target database to translate into
     * @return accessions from the target database
     * @throws IOException
     */
    public Set<String> getAccession(String database, String accession, String targetdatabase) throws IOException {
        final Set<String> accessions = new HashSet<String>();

        BooleanQuery bq = new BooleanQuery();
        if (database != null)
            bq.add(new TermQuery(new Term("gp_database", database)), BooleanClause.Occur.MUST);
        bq.add(new TermQuery(new Term("gp_value", accession)), BooleanClause.Occur.MUST);
        bq.add(new TermQuery(new Term("gp_value", accession)), BooleanClause.Occur.MUST);
        if (targetdatabase != null)
            bq.add(new TermQuery(new Term("t_database", targetdatabase)), BooleanClause.Occur.MUST);

        final Collector docCollector = new Collector() {

            private int docBase;

            @Override
            public void setScorer(Scorer scorer) throws IOException {
            }

            @Override
            public void collect(int doc) throws IOException {
                try {
                    accessions.add(indexSearcher.doc(doc + docBase).get("t_value"));
                } catch (Exception e) {
                    log.error(e);
                    e.printStackTrace();
                }
            }

            @Override
            public void setNextReader(IndexReader reader, int docBase) throws IOException {
                this.docBase = docBase;
            }

            @Override
            public boolean acceptsDocsOutOfOrder() {
                return true;
            }
        };

        indexSearcher.search(bq, docCollector);

        return accessions;
    }

}
